{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/rchaves/Projects/langwatch-saas/langwatch/langwatch_nlp/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import langwatch\n",
    "import dotenv\n",
    "\n",
    "dotenv.load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "from enum import Enum\n",
      "\n",
      "\n",
      "class Model(Enum):\n",
      "    high = 'high'\n",
      "    medium = 'medium'\n",
      "    low = 'low'\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<enum 'Model'>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datamodel_code_generator import InputFileType, generate\n",
    "import json\n",
    "import re\n",
    "import io\n",
    "from contextlib import redirect_stdout\n",
    "\n",
    "json_schema = {\n",
    "    \"type\": \"enum\",\n",
    "    \"enum\": [\"high\", \"medium\", \"low\"],\n",
    "    # \"properties\": {\n",
    "    #     \"importance\": {\"type\": \"number\", \"title\": \"Importance\"},\n",
    "    #     \"reasoning\": {\"type\": \"string\", \"title\": \"Reasoning\"},\n",
    "    # },\n",
    "    # # \"required\": [\"reasoning\", \"importance\"],\n",
    "    # \"title\": \"NestedAnswerImportance\",\n",
    "}\n",
    "\n",
    "# Generate code. The result is a string of the Python code.\n",
    "code_buffer = io.StringIO()\n",
    "\n",
    "model_name = json_schema.get(\"title\", \"Model\") # <- field name\n",
    "\n",
    "# Redirect stdout to the buffer while calling generate()\n",
    "with redirect_stdout(code_buffer):\n",
    "    generate(\n",
    "        json.dumps(json_schema),\n",
    "        input_file_type=InputFileType.JsonSchema,\n",
    "        class_name=model_name\n",
    "    )\n",
    "\n",
    "# Get the generated code as a string\n",
    "output = code_buffer.getvalue()\n",
    "output = re.sub(r\"# generated by[\\s\\S]*?from __future__ import annotations\", \"\", output).strip()\n",
    "output = re.sub(r\"class (.*)?\\(BaseModel\\):\\n    __root__: \", r\"\\1 = \", output).strip()\n",
    "\n",
    "print(output)\n",
    "\n",
    "namespace = {}\n",
    "exec(output, namespace, namespace)\n",
    "Model = namespace[model_name]\n",
    "\n",
    "Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any, Optional, Type, cast\n",
    "import dspy\n",
    "from dspy.signatures.signature import Signature, _default_instructions\n",
    "from dspy.adapters.types.image import try_expand_image_tags\n",
    "\n",
    "from langwatch_nlp.studio.utils import SerializableWithStringFallback\n",
    "import pydantic\n",
    "from pydantic import Field\n",
    "from dspy.utils.callback import BaseCallback\n",
    "from dspy.adapters.chat_adapter import ChatAdapter\n",
    "from dspy.adapters.json_adapter import _get_structured_outputs_response_format\n",
    "import json\n",
    "\n",
    "\n",
    "class TemplateAdapter(dspy.JSONAdapter):\n",
    "    \"\"\"\n",
    "    This is a \"TemplateAdapter\" DSPy Adapter, that avoid modifying the messages as much as possible,\n",
    "    and instead uses a {{mustache}} template formating to fill in the inputs on the messages.\n",
    "\n",
    "    This adapter does not append any text to the system prompt like DSPy normally does and uses json for the outputs\n",
    "    by default, this matches much better what users expect comming from OpenAI standards, and will allow them to simply\n",
    "    pick up the same prompts and json schemas and use in any other frameworks as is, since all of them adhere to the\n",
    "    raw OpenAI way of interating with LLMs.\n",
    "    \"\"\"\n",
    "\n",
    "    def __call__(\n",
    "        self,\n",
    "        lm,\n",
    "        lm_kwargs: dict[str, Any],\n",
    "        signature: Type[Signature],\n",
    "        demos: list[dict[str, Any]],\n",
    "        inputs: dict[str, Any],\n",
    "    ) -> list[dict[str, Any]]:\n",
    "        # If the signature has only one output field and it's a string, we can use the text only completion\n",
    "        if self._use_text_only_completion(signature, inputs):\n",
    "            return ChatAdapter.__call__(self, lm, lm_kwargs, signature, demos, inputs)  # type: ignore\n",
    "\n",
    "        # Replace the DSPyProgramOutputs title from the json schema with the signature name to bias the LLM in the right direction instead of randomly towards DSPy\n",
    "        model = _get_structured_outputs_response_format(signature)\n",
    "        schema = model.model_json_schema()\n",
    "        if schema.get(\"title\", None) == \"DSPyProgramOutputs\":\n",
    "            new_name = signature.__name__.replace(\"Signature\", \"\")\n",
    "            schema[\"title\"] = new_name\n",
    "            model.__name__ = new_name\n",
    "            model.model_json_schema = lambda *args, **kwargs: schema\n",
    "        lm_kwargs[\"response_format\"] = model\n",
    "        return ChatAdapter.__call__(self, lm, lm_kwargs, signature, demos, inputs)  # type: ignore\n",
    "\n",
    "    def format(\n",
    "        self,\n",
    "        signature: Type[Signature],\n",
    "        demos: list[dict[str, Any]],\n",
    "        inputs: dict[str, Any],\n",
    "    ) -> list[dict[str, Any]]:\n",
    "        inputs_copy = dict(inputs)\n",
    "\n",
    "        # If the signature and inputs have conversation history, we need to format the conversation history and\n",
    "        # remove the history field from the signature.\n",
    "        history_field_name = cast(str, self._get_history_field_name(signature))\n",
    "        if history_field_name:\n",
    "            # In order to format the conversation history, we need to remove the history field from the signature.\n",
    "            signature_without_history = signature.delete(history_field_name)\n",
    "            conversation_history = self.format_conversation_history(\n",
    "                signature_without_history,\n",
    "                history_field_name,\n",
    "                inputs_copy,\n",
    "            )\n",
    "\n",
    "        _messages = getattr(signature, \"_messages\", Field(default=[])).default\n",
    "\n",
    "        instructions = signature.instructions\n",
    "        if instructions == _default_instructions(signature):\n",
    "            instructions = \"\"\n",
    "        messages = []\n",
    "        messages.append(\n",
    "            {\n",
    "                \"role\": \"system\",\n",
    "                \"content\": self._format_template_inputs(\n",
    "                    instructions, inputs_copy\n",
    "                ),\n",
    "            }\n",
    "        )\n",
    "        messages.extend(self.format_demos(signature, demos))\n",
    "        if history_field_name:\n",
    "            messages.extend(conversation_history)\n",
    "        messages.extend(\n",
    "            [\n",
    "                m | {\"content\": self._format_template_inputs(m[\"content\"], inputs_copy)}\n",
    "                for m in _messages\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        messages = try_expand_image_tags(messages)\n",
    "\n",
    "        return messages\n",
    "\n",
    "    def _format_template_inputs(\n",
    "        self, template: str, inputs: dict[str, Any]\n",
    "    ) -> dict[str, Any]:\n",
    "        \"\"\"\n",
    "        Format the template inputs filling the {{ input }} placeholders.\n",
    "        \"\"\"\n",
    "\n",
    "        class SafeDict(dict):\n",
    "            def __missing__(self, key):\n",
    "                return \"{{\" + key + \"}}\"\n",
    "\n",
    "        template_fmt = template.replace(\"{{\", \"{\").replace(\"}}\", \"}\")\n",
    "        str_inputs: dict[str, str] = {}\n",
    "        for k, v in inputs.items():\n",
    "            str_inputs[k] = (\n",
    "                v\n",
    "                if type(v) == str\n",
    "                else json.dumps(v, cls=SerializableWithStringFallback)\n",
    "            )\n",
    "        return template_fmt.format_map(SafeDict(str_inputs))  # type: ignore\n",
    "\n",
    "    def parse(self, signature, completion):\n",
    "        if len(signature.output_fields) == 0:\n",
    "            return {}\n",
    "\n",
    "        first_field = list(signature.output_fields.items())[0]\n",
    "        if self._use_text_only_completion(signature, completion):\n",
    "            return {first_field[0]: completion}\n",
    "\n",
    "        return super().parse(signature, completion)\n",
    "\n",
    "    def _use_text_only_completion(self, signature, completion):\n",
    "        return len(signature.output_fields) == 0 or (\n",
    "            len(signature.output_fields) == 1\n",
    "            and list(signature.output_fields.values())[0].annotation == str\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Prediction(\n",
       "    answer='The capital of France is Paris.'\n",
       ")"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import dspy\n",
    "from dspy.clients.lm import litellm\n",
    "\n",
    "with langwatch.trace() as trace:\n",
    "    trace.autotrack_dspy()\n",
    "\n",
    "    lm = dspy.LM(model=\"openai/gpt-4o-mini\")\n",
    "    dspy.configure(adapter=TemplateAdapter())\n",
    "\n",
    "    class AnswerSignature(dspy.Signature):\n",
    "        _messages = [{\n",
    "            \"role\": \"user\",\n",
    "            \"content\": \"User is asking: {{question}}\"\n",
    "        }]\n",
    "\n",
    "        question: str = dspy.InputField(description=\"The question to answer\")\n",
    "        answer: str = dspy.OutputField(description=\"The answer to the question\")\n",
    "\n",
    "    predict = dspy.Predict(AnswerSignature)\n",
    "    predict.set_lm(lm)\n",
    "\n",
    "    prediction = predict(question=\"what is the capital of France??????????/?/\")\n",
    "\n",
    "prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Prediction(\n",
       "    importance=<Model.high: 'high'>,\n",
       "    answer='The capital of France is Paris.'\n",
       ")"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import dspy\n",
    "from dspy.clients.lm import litellm\n",
    "\n",
    "with langwatch.trace() as trace:\n",
    "    trace.autotrack_dspy()\n",
    "\n",
    "    lm = dspy.LM(model=\"openai/gpt-4o-mini\")\n",
    "    dspy.configure(adapter=TemplateAdapter())\n",
    "\n",
    "\n",
    "    class AnswerImportance(dspy.Signature):\n",
    "        \"\"\"You are a helpful assistant that answers questions and provides information.\"\"\"\n",
    "\n",
    "        _messages = [{\n",
    "            \"role\": \"user\",\n",
    "            \"content\": \"User is asking: {{question}}\"\n",
    "        }]\n",
    "\n",
    "        question: str = dspy.InputField(description=\"The question to answer\")\n",
    "        importance: Model = dspy.OutputField(\n",
    "            description=\"The importance of the answer\"\n",
    "        )\n",
    "        answer: str = dspy.OutputField(description=\"The answer to the question\")\n",
    "\n",
    "    predict = dspy.Predict(AnswerImportance)\n",
    "    predict.set_lm(lm)\n",
    "\n",
    "    prediction = predict(question=\"what is the capital of France??????????/?/\")\n",
    "\n",
    "prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ModelResponse(id='chatcmpl-BT6f6U4scevsik1gXQJqN7Z9IaQoe', created=1746276820, model='gpt-4o-mini-2024-07-18', object='chat.completion', system_fingerprint='fp_0392822090', choices=[Choices(finish_reason='stop', index=0, message=Message(content='{\"importance\":\"high\",\"answer\":\"The capital of France is Paris.\"}', role='assistant', tool_calls=None, function_call=None, provider_specific_fields={'refusal': None}, annotations=[]))], usage=Usage(completion_tokens=16, prompt_tokens=292, total_tokens=308, completion_tokens_details=CompletionTokensDetailsWrapper(accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0, text_tokens=None), prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=0, cached_tokens=0, text_tokens=None, image_tokens=None)), service_tier='default')"
      ]
     },
     "execution_count": 73,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import langwatch\n",
    "import litellm\n",
    "\n",
    "with langwatch.trace() as trace:\n",
    "    trace.autotrack_litellm_calls(litellm)\n",
    "\n",
    "    litellm.completion(\n",
    "        model=\"openai/gpt-4o-mini\",\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"system\",\n",
    "                \"content\": \"\"\"Your input fields are:\n",
    "    1. `question` (str): The question to answer\n",
    "    Your output fields are:\n",
    "    1. `importance` (Model): The importance of the answer\n",
    "    2. `answer` (str): The answer to the question\n",
    "    All interactions will be structured in the following way, with the appropriate values filled in.\n",
    "\n",
    "    Inputs will have the following structure:\n",
    "\n",
    "    [[ ## question ## ]]\n",
    "    {question}\n",
    "\n",
    "    Outputs will be a JSON object with the following fields.\n",
    "\n",
    "    [[ ## importance ## ]]\n",
    "    {importance}        # note: the value you produce must adhere to the JSON schema: {\"type\": \"string\", \"title\": \"Model\"}\n",
    "\n",
    "    [[ ## answer ## ]]\n",
    "    {answer}\n",
    "    In adhering to this structure, your objective is:\n",
    "            Given the fields `question`, produce the fields `importance`, `answer`.\"\"\",\n",
    "            },\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": \"\"\"[[ ## question ## ]]\n",
    "    What is the capital of France????????????????\n",
    "\n",
    "    Respond with a JSON object in the following order of fields: `importance` (must be formatted as a valid Python Model), then `answer`.\"\"\",\n",
    "            },\n",
    "            # {\"role\": \"user\", \"content\": \"What is the capital of France?\"},\n",
    "        ],\n",
    "        # response_format={\n",
    "        #     \"type\": \"json_schema\",\n",
    "        #     \"json_schema\": {\n",
    "        #         \"name\": \"DSPyProgramOutputs\",\n",
    "        #         \"strict\": True,\n",
    "        #         \"schema\": {\n",
    "        #             \"additionalProperties\": False,\n",
    "        #             \"$defs\": {\"Model\": {\"title\": \"Model\", \"type\": \"string\"}},\n",
    "        #             \"properties\": {\n",
    "        #                 # \"importance\": {\"title\": \"Importance\", \"type\": \"number\"},\n",
    "        #                 # \"importance\": {\"$ref\": \"#/$defs/Model\"},\n",
    "        #                 \"importance\": {\"title\": \"Importance\", \"type\": \"string\", \"enum\": [\"high\", \"medium\", \"low\"]},\n",
    "        #                 \"answer\": {\"title\": \"Answer\", \"type\": \"string\"},\n",
    "        #             },\n",
    "        #             \"required\": [\"importance\", \"answer\"],\n",
    "        #             \"title\": \"DSPyProgramOutputs\",\n",
    "        #             \"type\": \"object\",\n",
    "        #         },\n",
    "        #     },\n",
    "        # },\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
